package com.mahiru.mahiruaiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.jdbc.core.JdbcTemplate;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * @Author Enkidu
 * @Date 2025/7/6 16:38
 */
// @Configuration
public class PgVectorVectorStoreConfig {

    @Resource
    private ChatAppDocumentLoader chatAppDocumentLoader;

    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashScopeEmbeddingModel) {
        PgVectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashScopeEmbeddingModel)
                // 这里的维度需要和嵌入模型的输出维度一致
                .dimensions(1536)
                // 使用余弦距离
                .distanceType(COSINE_DISTANCE)
                // 使用HNSW索引
                .indexType(HNSW)
                // 如果表不存在则创建
                .initializeSchema(true)
                // 指定模式名称
                .schemaName("public")
                // 指定向量表名称
                .vectorTableName("chat_app_vector_store")
                // 每次批量插入的最大文档数量
                .maxDocumentBatchSize(10000)
                .build();
        // List<Document> documentList = chatAppDocumentLoader.loadMarkdowns();
        // vectorStore.add(documentList);
        return vectorStore;
    }
}
